package yanglib.network.nio;

import yanglib.membership.Membership;
import yanglib.message.Message;
import yanglib.message.MessageDeserializer;
import yanglib.network.MessageHandler;
import yanglib.network.Network;

import java.io.IOException;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.*;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.Set;
import java.util.concurrent.LinkedBlockingQueue;

/**
 * Created with IntelliJ IDEA.
 * User: user
 * Date: 9/15/14
 * Time: 2:41 PM
 * To change this template use File | Settings | File Templates.
 */
public class NIONetwork extends Network {

    private static final int MAX_MSG_SIZE = 8192;

    private ServerSocketChannel ss;
    private Selector selector;
    private HashMap<Integer, Session> outgoingSessions = new HashMap<Integer, Session>();

    final ByteBuffer directBuffer = ByteBuffer.allocateDirect(64 * 1024);

    private class Session {

        private SocketChannel sock;
        private SelectionKey sk;

        private ByteBuffer lenBuffer = ByteBuffer.allocate(4);
        ByteBuffer incomingBuffer = lenBuffer;

        LinkedBlockingQueue<ByteBuffer> outgoingBuffers = new LinkedBlockingQueue<ByteBuffer>();

        private int target = -1; //outgoing session has target ID. incoming session does not.

        public Session(SocketChannel sc, SelectionKey sk) {
            this.sock = sc;
            this.sk = sk;
        }

        public void setTarget(int target){
            this.target = target;
        }

        public int getTarget(){
            return this.target;
        }

        public SocketChannel getSocketChannel(){
            return this.sock;
        }

        public void sendBuffer(ByteBuffer bb) {
            try {
                synchronized (this) {
                    sk.selector().wakeup();
                    outgoingBuffers.add(bb);
                    if (sk.isValid()) {
                        sk.interestOps(sk.interestOps() | SelectionKey.OP_WRITE);
                    }
                }
            } catch (Exception e) {
                e.printStackTrace();
            }
        }

        private boolean readLength(SelectionKey k) throws IOException {
            //System.out.println("readLength");
            // Read the length, now get the buffer
            int len = lenBuffer.getInt();
            //System.out.println("readLength="+len);
            incomingBuffer = ByteBuffer.allocate(len);
            return true;
        }

        private void readPayload() throws IOException, InterruptedException {
            if (incomingBuffer.remaining() != 0) { // have we read length bytes?
                int rc = sock.read(incomingBuffer); // sock is non-blocking, so ok
                if (rc < 0) {
                    throw new IOException(
                            "Unable to read additional data from client "
                                    + sock);
                }
            }

            if (incomingBuffer.remaining() == 0) { // have we read length bytes?
                Message msg = msgDeserializer.deserialize(incomingBuffer.array(), 0, incomingBuffer.position());
                //System.out.println("handling "+msg);
                msgHandler.handleMessage(msg);
                lenBuffer.clear();
                incomingBuffer = lenBuffer;
            }
        }

        public void doIO(SelectionKey k) throws IOException {
            //System.out.println("doIO");
            try {
                if (k.isReadable()) {
                    //System.out.println("doIO read");
                    int rc = sock.read(incomingBuffer);
                    if (rc < 0) {
                        throw new IOException(
                                "Unable to read additional data from "+sock);
                    }
                    if (incomingBuffer.remaining() == 0) {
                        boolean isPayload;
                        if (incomingBuffer == lenBuffer) { // start of next request
                            incomingBuffer.flip();
                            isPayload = readLength(k);
                            incomingBuffer.clear();
                        } else {
                            // continuation
                            isPayload = true;
                        }
                        if (isPayload) { // not the case for 4letterword
                            readPayload();
                        } else {
                            // four letter words take care
                            // need not do anything else
                            return;
                        }
                    }
                }
                if (k.isWritable()) {
                    //System.out.println("doIO write");
                    // ZooLog.logTraceMessage(LOG,
                    // ZooLog.CLIENT_DATA_PACKET_TRACE_MASK
                    // "outgoingBuffers.size() = " +
                    // outgoingBuffers.size());
                    if (outgoingBuffers.size() > 0) {
                        // ZooLog.logTraceMessage(LOG,
                        // ZooLog.CLIENT_DATA_PACKET_TRACE_MASK,
                        // "sk " + k + " is valid: " +
                        // k.isValid());

                    /*
                          * This is going to reset the buffer position to 0 and the
                          * limit to the size of the buffer, so that we can fill it
                          * with data from the non-direct buffers that we need to
                          * send.
                          */
                        directBuffer.clear();

                        for (ByteBuffer b : outgoingBuffers) {
                            if (directBuffer.remaining() < b.remaining()) {
                            /*
                                    * When we call put later, if the directBuffer is to
                                    * small to hold everything, nothing will be copied,
                                    * so we've got to slice the buffer if it's too big.
                                    */
                                b = (ByteBuffer) b.slice().limit(
                                        directBuffer.remaining());
                            }
                        /*
                               * put() is going to modify the positions of both
                               * buffers, put we don't want to change the position of
                               * the source buffers (we'll do that after the send, if
                               * needed), so we save and reset the position after the
                               * copy
                               */
                            int p = b.position();
                            directBuffer.put(b);
                            b.position(p);
                            if (directBuffer.remaining() == 0) {
                                break;
                            }
                        }
                    /*
                          * Do the flip: limit becomes position, position gets set to
                          * 0. This sets us up for the write.
                          */
                        directBuffer.flip();

                        int sent = sock.write(directBuffer);
                        ByteBuffer bb;

                        // Remove the buffers that we have sent
                        while (outgoingBuffers.size() > 0) {
                            bb = outgoingBuffers.peek();

                            int left = bb.remaining() - sent;
                            if (left > 0) {
                            /*
                                    * We only partially sent this buffer, so we update
                                    * the position and exit the loop.
                                    */
                                bb.position(bb.position() + sent);
                                break;
                            }

                        /* We've sent the whole buffer, so drop the buffer */
                            sent -= bb.remaining();
                            outgoingBuffers.remove();
                        }
                        // ZooLog.logTraceMessage(LOG,
                        // ZooLog.CLIENT_DATA_PACKET_TRACE_MASK, "after send,
                        // outgoingBuffers.size() = " + outgoingBuffers.size());
                    }

                    synchronized (this) {
                        if (outgoingBuffers.size() == 0) {
                            sk.interestOps(sk.interestOps()
                                    & (~SelectionKey.OP_WRITE));
                        } else {
                            sk
                                    .interestOps(sk.interestOps()
                                            | SelectionKey.OP_WRITE);
                        }
                    }
                }
            } catch (CancelledKeyException e) {
               e.printStackTrace();
            } catch (IOException e) {
                e.printStackTrace();
            } catch(InterruptedException e){
                e.printStackTrace();
            }
            //System.out.println("doIO "+(System.currentTimeMillis()-start));
        }
    }

    public NIONetwork(Membership membership, MessageHandler msgHandler, MessageDeserializer msgDeserializer, int selfID) {
        super(membership, msgHandler, msgDeserializer, selfID);
        int port = membership.getAddress(selfID).getPort();
        System.out.println("binding to "+membership.getAddress(selfID));
        try {
            selector = Selector.open();

            this.ss = ServerSocketChannel.open();
            ss.socket().setReuseAddress(true);
            ss.socket().bind(membership.getAddress(selfID));
            ss.configureBlocking(false);
            ss.register(selector, SelectionKey.OP_ACCEPT);
            new MainThread().start();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    private class MainThread extends Thread{
        public void run() {
            while (!ss.socket().isClosed()) {
                Session s = null;
                try {
                    selector.select(1000);
                    Set<SelectionKey> selected;
                    synchronized (this) {
                        selected = selector.selectedKeys();
                    }
                   // if(selected.size()!=0)
                   //     System.out.println(selected);
                    ArrayList<SelectionKey> selectedList = new ArrayList<SelectionKey>(
                            selected);
                    Collections.shuffle(selectedList);
                    for (SelectionKey k : selectedList) {
                        if ((k.readyOps() & SelectionKey.OP_ACCEPT) != 0) {
                            System.out.println("accept event");
                            SocketChannel sc = ((ServerSocketChannel) k
                                    .channel()).accept();
                            InetAddress ia = sc.socket().getInetAddress();
                           // System.out.println("Accepted socket connection from "
                            //        + sc.socket().getRemoteSocketAddress());
                            sc.configureBlocking(false);
                            SelectionKey sk = sc.register(selector,
                                    SelectionKey.OP_READ);
                            Session cnxn = new Session(sc, sk);
                            sk.attach(cnxn);
                        } else if ((k.readyOps() & (SelectionKey.OP_READ | SelectionKey.OP_WRITE)) != 0) {

                            //System.out.println("read/write event");
                            s = (Session) k.attachment();
                            if(s.getSocketChannel().isConnected())
                                s.doIO(k);
                        } else if((k.readyOps() & SelectionKey.OP_CONNECT) != 0) {

                            SocketChannel sc = ((SocketChannel) k.channel());
                            System.out.println("connect event" + sc);
                            s = (Session) k.attachment();
                            sc.finishConnect();
                            k.interestOps(k.interestOps() | SelectionKey.OP_READ);
                            k.interestOps(k.interestOps() | SelectionKey.OP_WRITE);

                        }
                        else{
                            System.out.println("Unexpected ops in select "
                                    + k.readyOps());
                        }
                    }
                    selected.clear();
                } catch (RuntimeException e) {
                    e.printStackTrace();
                    removeSession(s);
                } catch (Exception e) {
                    e.printStackTrace();
                    removeSession(s);
                }
            }
            //clear();
            System.out.println("NIOServerCnxn factory exited run method");
        }
    }

    private synchronized void removeSession(Session s){
        if(s!= null && s.getTarget() != -1){
            System.out.println("Remove session "+s.getTarget());
            this.outgoingSessions.remove(s.getTarget());
        }
    }

    @Override
    protected void send(byte[] bytes, int offset, int len, InetSocketAddress target) throws IOException {
        //To change body of implemented methods use File | Settings | File Templates.
    }

    @Override
    public synchronized void send(Message message, int target) throws IOException {
        byte []bytes = new byte[MAX_MSG_SIZE];
        ByteBuffer buffer = ByteBuffer.wrap(bytes);
        buffer.putInt(0); //place holder for length
        message.serialize(buffer);
        int size = buffer.position() - 4;
        buffer.flip();
        buffer.putInt(size);
        buffer.rewind();
        if(outgoingSessions.containsKey(target)){
            outgoingSessions.get(target).sendBuffer(buffer);
        }
        else{
            SocketChannel sock;
            sock = SocketChannel.open();
            sock.configureBlocking(false);
            sock.socket().setSoLinger(false, -1);
            sock.socket().setTcpNoDelay(true);
            SelectionKey sockKey = sock.register(selector, SelectionKey.OP_CONNECT);
            Session session = new Session(sock, sockKey);
            sockKey.attach(session);
            System.out.println("connect to "+target+" "+membership.getAddress(target));

            sock.connect(membership.getAddress(target));
            session.setTarget(target);
            outgoingSessions.put(target, session);
            session.sendBuffer(buffer);
        }
    }
}
